import os
import numpy as np
import torch
from PIL import Image


class MasksForUnet(torch.utils.data.Dataset):
    def __init__(self, images_root, masks_root, transforms=None):
        self.images_root = images_root
        self.masks_root = masks_root
        self.transforms = transforms

        self.imgs = list(sorted(os.listdir(images_root)))
        self.masks = list(sorted(os.listdir(masks_root)))

    def __getitem__(self, idx):
        img_path = os.path.join(self.images_root, self.imgs[idx])
        mask_path = os.path.join(self.masks_root, self.masks[idx])

        img = Image.open(img_path).convert("RGB")
        img = np.array(img)

        target = Image.open(mask_path)
        target = np.array(target)

        if self.transforms is not None:
            img = self.transforms(img)
            target = self.transforms(target)

        return img, target

    def get_categories(self):
        return None

    def __len__(self):
        return len(self.imgs)
